# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
"""
Train and eval functions used in main.py
Modified from: https://github.com/facebookresearch/deit
"""
import math
import sys
import os
from typing import Iterable, Optional

import torch
import torch.nn.functional as F

from timm.data import Mixup
from timm.utils import accuracy, ModelEma

from anns.tost.vis.losses import DistillationLoss
import anns.tost.vis.utils as utils


def train_one_epoch(model: torch.nn.Module, criterion: DistillationLoss,
                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
                    model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None,
                    set_training_mode=True, surgery=None):
    model.train(set_training_mode)

    if surgery:
        model.module.patch_embed.eval()

    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    header = 'Epoch: [{}]'.format(epoch)
    print_freq = 10

    for batch in metric_logger.log_every(data_loader, print_freq, header):
        samples, targets = batch[0], batch[1]

        samples = samples.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)

        if mixup_fn is not None:
            samples, targets = mixup_fn(samples, targets)

        with torch.cuda.amp.autocast():
            outputs = model(samples)

            loss = criterion(samples, outputs, targets)

        loss_value = loss.item()

        if not math.isfinite(loss_value):
            print("Loss is {}, stopping training".format(loss_value))
            sys.exit(1)

        optimizer.zero_grad()

        # this attribute is added by timm on one optimizer (adahessian)
        is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
        loss_scaler(loss, optimizer, clip_grad=max_norm,
                    parameters=model.parameters(), create_graph=is_second_order)

        torch.cuda.synchronize()
        if model_ema is not None:
            model_ema.update(model)

        metric_logger.update(loss=loss_value)
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}


@torch.no_grad()
def evaluate(data_loader, model, device):
    criterion = torch.nn.CrossEntropyLoss()

    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Test:'

    # switch to evaluation mode
    model.eval()

    for batch in metric_logger.log_every(data_loader, 10, header):
        images, target = batch[0], batch[1]

        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)

        # compute output
        with torch.cuda.amp.autocast():
            output = model(images)

        loss = criterion(output, target)
        acc1, acc5 = accuracy(output, target, topk=(1, 5))

        batch_size = images.shape[0]
        metric_logger.update(loss=loss.item())
        metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
        metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
          .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))

    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}



def save_sample_fn(x:torch.Tensor, sample_idx:int, ds_folder_:str = './work/data/exp1') -> str:
    ds_folder = ds_folder_
    raw_path = f'{sample_idx:012d}'
    for pos_idx in range(0, len(raw_path)-2, 2):
        ds_folder = f'{ds_folder}/d{raw_path[pos_idx: pos_idx+2]}'
        if not os.path.exists(ds_folder):
            os.mkdir(ds_folder)
    ds_fn = f'{ds_folder}/{raw_path}.pt'
    torch.save(x, ds_fn)
    return ds_fn

# python main.py --predict --model tost_cab_tiny_12_p16 --input-size 224 --full_crop --pretrained ../work/outputs/tiny_20250927/best_model.pth
@torch.no_grad()
def predict(args, dataset_train, sampler_train, model, device):
    print(f'开始进行推理')
    data_loader = torch.utils.data.DataLoader(
        dataset_train, #sampler=sampler_train,
        batch_size=1,
        num_workers=1,
        pin_memory=args.pin_mem,
        drop_last=True,
    )
    criterion = torch.nn.CrossEntropyLoss()
    metric_logger = utils.MetricLogger(delimiter="  ")
    header = '统计:'
    # switch to evaluation mode
    prob_oks = [] # 正确记录概率列表
    prob_errors = [] # 错误记录要概率列表
    model.eval()
    sample_idx = 0
    ds_folder = './work/data/exp6'
    ok_folder = f'{ds_folder}/oks'
    error_folder = f'{ds_folder}/errors'
    ds_dict = {}
    percents = {}
    err_percs = {}
    ok_num, error_num = 0, 0
    for i in range(0, 101, 1):
        percents[i] = 0
        err_percs[i] = 0
    # i_debug = 1
    for batch in metric_logger.log_every(data_loader, 100, header):
        images, target = batch[0], batch[1]
        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)
        # if i_debug == 1:
        #     X = images.detach().cpu().squeeze(0)
        #     # X = X.permute(1, 2, 0)
        #     X = X / 2 + 0.5
        #     # if X.max() < 1.0:
        #     #     X = X*255
        #     img_raw = X.numpy()
        #     plt.imshow(np.transpose(img_raw, (1, 2, 0)))
        #     plt.show()
        #     return
        # compute output
        with torch.amp.autocast(device_type='cuda'):
            output = model(images)
        v = 1 / 0
        probs = F.softmax(output, dim=1)
        max_val, max_idx = torch.max(probs, dim=1)
        max_val_f = max_val.detach().cpu().item()
        max_val_perc = int(max_val_f*100)
        max_idx_int = max_idx.detach().cpu().item()
        target_int = target.detach().cpu().item()
        ds_fn = ''
        if max_idx_int == target_int:
            ok_num += 1
            prob_oks.append(max_val_f)
            max_val_int = int(max_val_perc)
            percents[max_val_int] += 1
            for img in images:
                img = img.detach().cpu()
                ds_fn = save_sample_fn(img, sample_idx, ok_folder)
                ds_dict[ds_fn] = target_int
        else:
            error_num += 1
            prob_errors.append(max_val_f)
            max_val_int = int(max_val_perc)
            err_percs[max_val_int] += 1
            for img in images:
                img = img.detach().cpu()
                ds_fn = save_sample_fn(img, sample_idx, error_folder)
                ds_dict[ds_fn] = target_int
        # print(f'probs: {max_val}; cls_id: {max_idx};')
        # print(f'target: {target.shape}; {target};')
        sample_idx += 1
        if sample_idx % 1000 == 0:
            print(f'总样本：{sample_idx}={ok_num} + {error_num};')
    prob_oks_t = torch.tensor(prob_oks)
    prob_errors_t = torch.tensor(prob_errors)
    ok_mean = torch.mean(prob_oks_t)
    ok_var = torch.var(prob_oks_t)
    ok_min = torch.min(prob_oks_t)
    error_mean = torch.mean(prob_errors_t)
    error_var = torch.var(prob_errors_t)
    error_max = torch.max(prob_errors_t)
    print(f'### 正确样本：均值：{ok_mean}; 方差：{ok_var}; 最小值：{ok_min}; 数量：{ok_num};')
    print(f'### 错误样本：均值：{error_mean}; 方差：{error_var}; 最大值：{error_max}; 数量：{error_num};')
    with open(f'{ds_folder}/raw.txt', 'w', encoding='utf-8') as wfd:
        for k,v in ds_dict.items():
            wfd.write(f'{k[1:]},{v}\n')
    x_t = []
    y_t = []
    for k, v in percents.items():
        x_t.append(k)
        y_t.append(v)
    plt.bar(x_t, y_t)
    plt.show()
    x_err = []
    y_err = []
    for k, v in err_percs.items():
        x_err.append(k)
        y_err.append(v)
    plt.bar(x_err, y_err)
    plt.show()
    print(f'### ^_^ The End! ^_^')
